import os


# rootpath = r'G:\ml2021\hw3\food-11'
# trainpath = os.path.join(rootpath, 'training/labeled')
# valpath = os.path.join(rootpath, 'validation/labeled')
#
# file_handle = open('train.txt', mode='w')
#
# folders = os.listdir(trainpath)
# for typefolder in folders:
#     fullpath = valpath = os.path.join(trainpath, typefolder)
#     files = os.listdir(fullpath)
#     for file in files:
#         file_handle.write(f'file {typefolder} \n')


def getlabel(input):
    """输入是文件夹名，输出为数字label"""
    return int(input)


def gene(rootpath, subpath, savename, getlabel=None):
    file_handle = open(savename, mode='w')
    folderpath = os.path.join(rootpath, subpath)
    folders = os.listdir(folderpath)
    for typefolder in folders:
        fullpath = os.path.join(folderpath, typefolder)
        files = os.listdir(fullpath)
        label = typefolder
        if getlabel is not None:
            label = getlabel(typefolder)
        for file in files:
            filepath_over_root = os.path.join(subpath, typefolder)
            filepath_over_root = os.path.join(filepath_over_root, file)
            filepath_over_root = filepath_over_root.__str__().replace("\\", '/')
            file_handle.write(f'{filepath_over_root} {label} \n')


if __name__ == '__main__':
    rootpath = r'G:\ml2021\hw3\food-11'
    trainpath = r'training\labeled'
    valpath = 'validation'
    gene(rootpath, valpath, 'val.txt', getlabel)
    gene(rootpath, trainpath, 'train.txt', getlabel)
